# -- coding: utf8 --
import os
import csv
import time
import glob
import pandas as pd
from PIL import Image

# label数据归一化
# size为图片宽高
# box为label的xmin/xmax/ymin/ymax
# x,y为归一化后label的中心坐标
# w,h为归一化后label的宽高
def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_(box):
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    return (x,y,w,h)

# 要下的数据集 train, test, validation
runMode = 'validation'
# 类别，输入要和class-descriptions-boxable.csv里的名称完全匹配
# 可一次下载多个类别
classes = ["Goldfish", "Fish"]
# 图片路径
img_path = 'JPEGImages/' + runMode + '/'
# txt保存路径
txt_path = 'labels/' + runMode + '/'

# 读取class-descriptions-boxable.csv的所有内容
with open('./Annotations/class-descriptions-boxable.csv', 'r') as infile:
    reader = csv.reader(infile)
    # 创建一个字典，键为类名，值为LabelName(class_ids)
    class_dict = {rows[1]: rows[0] for rows in reader}

# 根据类别读取LabelName
class_ids = [] # 存储LabelName，与classes长度相同
label_dict = {} # LabelName与索引关联的字典
for your_class in classes:
    if your_class in class_dict:
        class_ids.append(class_dict[your_class])
        label_dict.update({class_dict[your_class]:classes.index(your_class)})

# 读取train-annotations-bbox.csv所有内容
df = pd.read_csv('./Annotations/' + runMode + '-annotations-bbox.csv')
# 根据LabelName获取classes对应Label的信息
# 去掉其他classes的label
df = df.loc[df['LabelName'].isin(class_ids)]
# 去掉边界框跨越一组对象的图片
df = df[df['IsGroupOf'] == 0]
# 将索引重置为自然序列
df = df.reset_index(drop=True)

# if os.path.isdir(txt_path):
#     os.system('rd ' + '/s/q ' + txt_path)
#     os.system('md ' + txt_path)
# else:
#     os.system('mkdir -p ' + txt_path)

# path_file=glob.glob(pathname=img_path + '/*.jpg')

n = 0
for index in range(df.shape[0]):
    for label in class_ids:
        if label == df['LabelName'].iloc[index] and os.path.isfile(img_path + df['ImageID'].iloc[index] + ".jpg"):
            # 加载图片
            img = Image.open(img_path + df['ImageID'].iloc[index] + ".jpg")
            # 获取图片宽高
            size = img.size
            # 获取label信息
            box = df[['XMin', 'XMax', 'YMin', 'YMax']].iloc[index].values
            label_info = convert_(box)
            n=n+1
            # 将label信息写入txt文件
            annotations = open(txt_path + df['ImageID'].iloc[index] + ".txt", 'a')
            annotations.write(str(label_dict[label]) + " " + " ".join([str(a) for a in label_info]) + '\n')
            annotations.close()
            print("(" + str(n) + "/" + str(df.shape[0]) + ") " + time.strftime('%Y-%m-%d %H:%M:%S\n', time.localtime(time.time())))
